"""
Authentication Manager - Enhanced support for multi-website authentication flows.

This module handles the authentication process for multiple websites,
including session state capture, cloud storage integration, and
browser extension backup functionality.
"""

import os
import sys
import json
import time
import zipfile
import shutil
import tempfile
import subprocess
import logging
from typing import Dict, List, Optional, Any, Tuple
from pathlib import Path
from datetime import datetime

from models import WebsiteProfile, UserSession, AuthMethod, SessionStatus
from config import upload_file_to_oss


class AuthenticationManager:
    """
    Manages authentication processes and state persistence across multiple websites.
    
    Handles website-specific login procedures, session state capture,
    cloud storage synchronization, and browser extension backup.
    """
    
    def __init__(self, temp_dir: str = "temp_auth"):
        """
        Initialize the Authentication Manager.
        
        Args:
            temp_dir: Directory for temporary authentication files
        """
        self.temp_dir = Path(temp_dir)
        self.temp_dir.mkdir(exist_ok=True)
        
        self.logger = self._setup_logging()
        self.chrome_path = self._find_chrome_executable()
        
        if not self.chrome_path:
            self.logger.error("Chrome browser not found")
            raise RuntimeError("Chrome browser is required but not found")
        
        self.logger.info("AuthenticationManager initialized")
    
    def _setup_logging(self) -> logging.Logger:
        """Setup logging configuration."""
        logger = logging.getLogger("AuthenticationManager")
        logger.setLevel(logging.INFO)
        
        if not logger.handlers:
            handler = logging.StreamHandler()
            formatter = logging.Formatter(
                '%(asctime)s - %(name)s - %(levelname)s - %(message)s'
            )
            handler.setFormatter(formatter)
            logger.addHandler(handler)
        
        return logger
    
    def _find_chrome_executable(self) -> Optional[str]:
        """Find Chrome executable on the system."""
        if sys.platform == "win32":
            possible_paths = [
                r"C:\Program Files\Google\Chrome\Application\chrome.exe",
                r"C:\Program Files (x86)\Google\Chrome\Application\chrome.exe",
                os.path.expanduser(r"~\AppData\Local\Google\Chrome\Application\chrome.exe"),
            ]
        elif sys.platform == "darwin":  # macOS
            possible_paths = [
                "/Applications/Google Chrome.app/Contents/MacOS/Google Chrome",
                os.path.expanduser("~/Applications/Google Chrome.app/Contents/MacOS/Google Chrome"),
            ]
        else:  # Linux
            possible_paths = [
                "/usr/bin/google-chrome",
                "/usr/bin/chromium-browser",
                "/snap/bin/chromium",
            ]
        
        for path in possible_paths:
            if os.path.isfile(path):
                self.logger.info(f"Found Chrome: {path}")
                return path
        
        # Try command line
        try:
            result = subprocess.run(["chrome", "--version"], 
                                   capture_output=True, text=True, timeout=5)
            if result.returncode == 0:
                self.logger.info("Found Chrome: chrome")
                return "chrome"
        except Exception:
            pass
        
        return None
    
    def _get_system_chrome_profile_dir(self) -> Optional[str]:
        """Get system Chrome default profile directory."""
        try:
            if sys.platform == "win32":
                profile_paths = [
                    os.path.expanduser(r"~\AppData\Local\Google\Chrome\User Data\Default"),
                    os.path.expanduser(r"~\AppData\Local\Google\Chrome\User Data\Profile 1"),
                ]
            elif sys.platform == "darwin":  # macOS
                profile_paths = [
                    os.path.expanduser("~/Library/Application Support/Google/Chrome/Default"),
                    os.path.expanduser("~/Library/Application Support/Google/Chrome/Profile 1"),
                ]
            else:  # Linux
                profile_paths = [
                    os.path.expanduser("~/.config/google-chrome/Default"),
                    os.path.expanduser("~/.config/google-chrome/Profile 1"),
                ]
            
            for path in profile_paths:
                if os.path.exists(path):
                    self.logger.debug(f"Found system Chrome profile: {path}")
                    return path
            
            return None
            
        except Exception as e:
            self.logger.error(f"Failed to get system Chrome profile: {e}")
            return None
    
    def _backup_extensions_data(self, profile_dir: str, username: str) -> bool:
        """
        Backup Chrome extensions data to cloud storage.
        
        Args:
            profile_dir: Chrome profile directory
            username: User identifier
            
        Returns:
            True if successful, False otherwise
        """
        try:
            self.logger.info("Backing up extensions data...")
            
            # Create temporary backup directory
            temp_backup_dir = self.temp_dir / f"extensions_backup_{username}_{int(time.time())}"
            temp_backup_dir.mkdir(exist_ok=True)
            
            try:
                # Extensions directories to backup
                extension_dirs = [
                    "Extensions",
                    "Extension State", 
                    "Extension Rules",
                    "Local Extension Settings",
                    "Sync Extension Settings",
                    "Local Storage",
                    "Session Storage", 
                    "IndexedDB"
                ]
                
                # Extensions files to backup
                extension_files = [
                    "Extension Cookies",
                    "Extension Cookies-journal",
                    "Web Data",
                    "Web Data-journal"
                ]
                
                copied_count = 0
                
                # Copy directories
                for item in extension_dirs:
                    source_path = Path(profile_dir) / item
                    target_path = temp_backup_dir / item
                    
                    if source_path.exists() and source_path.is_dir():
                        try:
                            shutil.copytree(source_path, target_path)
                            copied_count += 1
                            self.logger.debug(f"Backed up directory: {item}")
                        except Exception as e:
                            self.logger.warning(f"Failed to backup directory {item}: {e}")
                
                # Copy files
                for item in extension_files:
                    source_path = Path(profile_dir) / item
                    target_path = temp_backup_dir / item
                    
                    if source_path.exists() and source_path.is_file():
                        try:
                            target_path.parent.mkdir(parents=True, exist_ok=True)
                            shutil.copy2(source_path, target_path)
                            copied_count += 1
                            self.logger.debug(f"Backed up file: {item}")
                        except Exception as e:
                            self.logger.warning(f"Failed to backup file {item}: {e}")
                
                if copied_count == 0:
                    self.logger.warning("No extension data found to backup")
                    return False
                
                # Create zip archive
                backup_filename = f"{username}_extensions_backup.zip"
                backup_path = self.temp_dir / backup_filename
                
                self.logger.info("Creating extensions backup archive...")
                with zipfile.ZipFile(backup_path, 'w', zipfile.ZIP_DEFLATED) as zipf:
                    for root, dirs, files in os.walk(temp_backup_dir):
                        for file in files:
                            file_path = Path(root) / file
                            arcname = file_path.relative_to(temp_backup_dir)
                            zipf.write(file_path, arcname)
                
                # Get archive size
                backup_size = backup_path.stat().st_size / (1024 * 1024)  # MB
                self.logger.info(f"Extensions backup size: {backup_size:.2f} MB")
                
                # Upload to cloud
                try:
                    self.logger.info("Uploading extensions backup to cloud...")
                    upload_file_to_oss(str(backup_path))
                    self.logger.info("Extensions backup uploaded successfully")
                    
                    # Clean up local file
                    backup_path.unlink()
                    return True
                    
                except Exception as e:
                    self.logger.error(f"Failed to upload extensions backup: {e}")
                    self.logger.info(f"Local backup saved: {backup_path}")
                    return False
                    
            finally:
                # Clean up temporary directory
                if temp_backup_dir.exists():
                    shutil.rmtree(temp_backup_dir)
                    
        except Exception as e:
            self.logger.error(f"Extensions backup failed: {e}")
            return False
    
    def authenticate_website(self, profile: WebsiteProfile, username: str, 
                           password: Optional[str] = None, 
                           backup_extensions: bool = True) -> Optional[UserSession]:
        """
        Authenticate user on specified website and capture session state.
        
        Args:
            profile: Website profile configuration
            username: User account identifier
            password: User password (optional for manual login)
            backup_extensions: Whether to backup browser extensions
            
        Returns:
            Created user session or None if failed
        """
        try:
            self.logger.info(f"Starting authentication for {profile.display_name}")
            
            # Determine authentication method
            if profile.auth_method == AuthMethod.MANUAL:
                return self._authenticate_manual(profile, username, backup_extensions)
            elif profile.auth_method == AuthMethod.FORM:
                return self._authenticate_form(profile, username, password, backup_extensions)
            elif profile.auth_method == AuthMethod.CUSTOM:
                return self._authenticate_custom(profile, username, password, backup_extensions)
            else:
                self.logger.error(f"Unsupported authentication method: {profile.auth_method}")
                return None
                
        except Exception as e:
            self.logger.error(f"Authentication failed: {e}")
            return None
    
    def _authenticate_manual(self, profile: WebsiteProfile, username: str, 
                           backup_extensions: bool) -> Optional[UserSession]:
        """Handle manual authentication flow."""
        try:
            # Determine profile directory
            system_profile = self._get_system_chrome_profile_dir()
            if backup_extensions and system_profile:
                profile_dir = system_profile
            else:
                # Create temporary profile
                timestamp = int(time.time())
                profile_dir = self.temp_dir / f"profile_{username}_{timestamp}"
                profile_dir.mkdir(exist_ok=True)
                profile_dir = str(profile_dir)
            
            # Launch Chrome for manual login
            chrome_args = [
                self.chrome_path,
                f"--user-data-dir={profile_dir}",
                "--start-maximized",
                profile.target_url
            ]
            
            # Add browser-specific arguments
            if profile.browser_config.extra_args:
                chrome_args.extend(profile.browser_config.extra_args)
            
            self.logger.info("Launching Chrome for manual authentication...")
            
            # Start Chrome process
            process = subprocess.Popen(
                chrome_args,
                stdout=subprocess.DEVNULL,
                stderr=subprocess.DEVNULL,
                stdin=subprocess.DEVNULL,
                creationflags=subprocess.CREATE_NEW_PROCESS_GROUP if sys.platform == 'win32' else 0
            )
            
            self.logger.info(f"Chrome process started (PID: {process.pid})")
            time.sleep(3)
            
            # Provide user instructions
            print("\n" + "="*60)
            print("📋 Please complete the following steps:")
            print("1. 🔐 Complete website login")
            print("2. 🔌 Install and configure required browser extensions")
            print("3. 🔑 Complete extension login and configuration")
            print("4. ✅ Verify all extensions are working correctly")
            print("5. ⌨️  Press Enter when complete to save session state")
            print("="*60)
            
            input("\nPress Enter to save login state...")
            
            # Request user to close Chrome
            print("📌 Please close Chrome browser, then press Enter to continue saving...")
            input()
            
            # Wait for process to terminate
            try:
                process.wait(timeout=5)
            except subprocess.TimeoutExpired:
                process.terminate()
                time.sleep(2)
            
            # Capture session state using Playwright
            auth_state = self._capture_session_state(profile_dir, profile.target_url)
            if not auth_state:
                return None
            
            # Backup extensions if requested
            extensions_backup_success = False
            if backup_extensions and system_profile:
                extensions_backup_success = self._backup_extensions_data(system_profile, username)
            
            # Create session
            session = self._create_session(profile, username, auth_state)
            
            # Upload session state to cloud
            try:
                auth_file = f"{username}_{profile.website_id}_auth_state.json"
                auth_path = self.temp_dir / auth_file
                
                with open(auth_path, 'w', encoding='utf-8') as f:
                    json.dump(auth_state, f, indent=2, ensure_ascii=False)
                
                upload_file_to_oss(str(auth_path))
                self.logger.info("Session state uploaded to cloud")
                auth_path.unlink()  # Clean up
                
            except Exception as e:
                self.logger.error(f"Failed to upload session state: {e}")
            
            if extensions_backup_success:
                self.logger.info("Extensions backup completed successfully")
            
            return session
            
        except Exception as e:
            self.logger.error(f"Manual authentication failed: {e}")
            return None
    
    def _authenticate_form(self, profile: WebsiteProfile, username: str, 
                          password: Optional[str], backup_extensions: bool) -> Optional[UserSession]:
        """Handle automated form-based authentication."""
        try:
            if not password:
                self.logger.error("Password required for form authentication")
                return None
            
            # Use Playwright for automated form submission
            from playwright.sync_api import sync_playwright
            
            with sync_playwright() as p:
                # Create temporary profile directory
                timestamp = int(time.time())
                profile_dir = self.temp_dir / f"form_auth_{username}_{timestamp}"
                
                # Launch persistent context
                context = p.chromium.launch_persistent_context(
                    user_data_dir=str(profile_dir),
                    headless=False,  # Show browser for user verification
                    channel="chrome",
                    viewport=profile.browser_config.viewport,
                    args=profile.browser_config.extra_args
                )
                
                page = context.new_page()
                
                # Navigate to login page
                login_url = profile.login_url or profile.target_url
                page.goto(login_url, timeout=30000)
                
                # Fill form fields
                if profile.selectors.username_field:
                    page.fill(profile.selectors.username_field, username)
                    time.sleep(1)
                
                if profile.selectors.password_field:
                    page.fill(profile.selectors.password_field, password)
                    time.sleep(1)
                
                # Handle captcha if present
                if profile.selectors.captcha_field:
                    print("🔍 Please solve the captcha and press Enter to continue...")
                    input()
                
                # Submit form
                if profile.selectors.submit_button:
                    page.click(profile.selectors.submit_button)
                else:
                    # Try common submit methods
                    try:
                        page.keyboard.press("Enter")
                    except Exception:
                        pass
                
                # Wait for navigation
                time.sleep(5)
                
                # Verify successful login
                if login_url != profile.target_url:
                    page.goto(profile.target_url, timeout=30000)
                    time.sleep(3)
                
                print("✅ Automated login completed. Please verify and press Enter...")
                input()
                
                # Capture session state
                auth_state_path = self.temp_dir / f"temp_auth_{username}_{int(time.time())}.json"
                context.storage_state(path=str(auth_state_path))
                
                # Read auth state
                with open(auth_state_path, 'r', encoding='utf-8') as f:
                    auth_state = json.load(f)
                
                auth_state_path.unlink()  # Clean up
                context.close()
                
                # Create session
                session = self._create_session(profile, username, auth_state)
                
                # Upload to cloud
                try:
                    auth_file = f"{username}_{profile.website_id}_auth_state.json"
                    auth_path = self.temp_dir / auth_file
                    
                    with open(auth_path, 'w', encoding='utf-8') as f:
                        json.dump(auth_state, f, indent=2, ensure_ascii=False)
                    
                    upload_file_to_oss(str(auth_path))
                    auth_path.unlink()
                    
                except Exception as e:
                    self.logger.error(f"Failed to upload session state: {e}")
                
                return session
                
        except ImportError:
            self.logger.error("Playwright not available for form authentication")
            return None
        except Exception as e:
            self.logger.error(f"Form authentication failed: {e}")
            return None
    
    def _authenticate_custom(self, profile: WebsiteProfile, username: str, 
                           password: Optional[str], backup_extensions: bool) -> Optional[UserSession]:
        """Handle custom script-based authentication."""
        try:
            # Custom authentication would execute user-defined scripts
            # This is a placeholder for custom automation scripts
            self.logger.info("Custom authentication not yet implemented")
            return None
            
        except Exception as e:
            self.logger.error(f"Custom authentication failed: {e}")
            return None
    
    def _capture_session_state(self, profile_dir: str, target_url: str) -> Optional[Dict[str, Any]]:
        """Capture session state using Playwright."""
        try:
            from playwright.sync_api import sync_playwright
            
            with sync_playwright() as p:
                # Connect to existing profile
                context = p.chromium.launch_persistent_context(
                    user_data_dir=profile_dir,
                    headless=True,
                    channel="chrome",
                    args=['--disable-web-security', '--no-first-run']
                )
                
                page = context.new_page()
                page.goto(target_url, timeout=30000)
                time.sleep(3)
                
                # Save state to temporary file
                temp_state_file = self.temp_dir / f"temp_state_{int(time.time())}.json"
                context.storage_state(path=str(temp_state_file))
                
                # Read state
                with open(temp_state_file, 'r', encoding='utf-8') as f:
                    auth_state = json.load(f)
                
                temp_state_file.unlink()  # Clean up
                context.close()
                
                # Validate state
                cookies_count = len(auth_state.get('cookies', []))
                if cookies_count > 0:
                    self.logger.info(f"Captured session state with {cookies_count} cookies")
                    return auth_state
                else:
                    self.logger.warning("No session data captured")
                    return None
                    
        except ImportError:
            self.logger.error("Playwright not available for session capture")
            return None
        except Exception as e:
            self.logger.error(f"Failed to capture session state: {e}")
            return None
    
    def _create_session(self, profile: WebsiteProfile, username: str, 
                       auth_state: Dict[str, Any]) -> UserSession:
        """Create UserSession object from authentication data."""
        from models import create_user_session
        
        session = create_user_session(profile.website_id, username, auth_state)
        session.status = SessionStatus.ACTIVE
        
        return session
    
    def download_session_state(self, username: str, website_id: str) -> Optional[Dict[str, Any]]:
        """
        Download session state from cloud storage.
        
        Args:
            username: User identifier
            website_id: Website identifier
            
        Returns:
            Session state dictionary or None if not found
        """
        try:
            import requests
            
            auth_file = f"{username}_{website_id}_auth_state.json"
            cloud_url = f"https://auth-share.oss-cn-shenzhen.aliyuncs.com/{auth_file}"
            
            self.logger.info(f"Downloading session state: {auth_file}")
            response = requests.get(cloud_url)
            response.raise_for_status()
            
            auth_state = response.json()
            self.logger.info("Session state downloaded successfully")
            return auth_state
            
        except Exception as e:
            self.logger.error(f"Failed to download session state: {e}")
            return None
    
    def download_extensions_backup(self, username: str, target_dir: str) -> bool:
        """
        Download and extract extensions backup from cloud storage.
        
        Args:
            username: User identifier
            target_dir: Directory to extract extensions data
            
        Returns:
            True if successful, False otherwise
        """
        try:
            import requests
            
            backup_file = f"{username}_extensions_backup.zip"
            cloud_url = f"https://auth-share.oss-cn-shenzhen.aliyuncs.com/{backup_file}"
            
            self.logger.info(f"Downloading extensions backup: {backup_file}")
            response = requests.get(cloud_url)
            response.raise_for_status()
            
            # Save backup file
            backup_path = self.temp_dir / backup_file
            with open(backup_path, "wb") as f:
                f.write(response.content)
            
            # Extract to target directory
            with zipfile.ZipFile(backup_path, 'r') as zipf:
                zipf.extractall(target_dir)
            
            # Clean up
            backup_path.unlink()
            
            self.logger.info("Extensions backup downloaded and extracted successfully")
            return True
            
        except Exception as e:
            self.logger.error(f"Failed to download extensions backup: {e}")
            return False
    
    def validate_session_state(self, auth_state: Dict[str, Any]) -> bool:
        """
        Validate session state data.
        
        Args:
            auth_state: Session state dictionary
            
        Returns:
            True if valid, False otherwise
        """
        try:
            if not isinstance(auth_state, dict):
                return False
            
            # Check required fields
            if 'cookies' not in auth_state:
                return False
            
            cookies = auth_state['cookies']
            if not isinstance(cookies, list):
                return False
            
            # Validate at least one cookie exists
            if len(cookies) == 0:
                return False
            
            # Validate cookie structure
            for cookie in cookies:
                if not isinstance(cookie, dict):
                    return False
                if not all(key in cookie for key in ['name', 'value', 'domain']):
                    return False
            
            return True
            
        except Exception:
            return False